package com.gearbox.platform.slurm.util;

import com.gearbox.core.constant.NodeStatusEnum;
import com.gearbox.core.constant.SlurmInstanceStatusEnum;
import com.gearbox.core.model.Job;
import com.gearbox.core.model.Node;
import com.gearbox.platform.slurm.exception.CmdExecuteResultParseException;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Locale;
import java.util.Optional;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

public class SlurmCmdExecuteResultParser {
    private static final int SINFO_STATE_INDEX = 3;

    private static final int SINFO_NODELIST_INDEX = 0;

    private static final int SINFO_MAX_COL = 4;

    private static final int SQUEUE_ID_INDEX = 0;

    private static final int SQUEUE_NAME_INDEX = 1;

    private static final int SQUEUE_USER_INDEX = 2;

    private static final int SQUEUE_CPUS_INDEX = 3;

    private static final int SQUEUE_MEMORY_INDEX = 4;

    private static final int SQUEUE_STATE_INDEX = 5;

    private static final int SQUEUE_MAX_COL = 6;

    private static final Pattern NODE_NAME_PREFIX_PATTERN = Pattern.compile("^[a-zA-Z0-9]+");

    private static final Pattern NODE_NAME_SUFFIX_SET_PATTERN = Pattern.compile("\\[[0-9]+-[0-9]+\\]$");

    private static final Pattern NODE_NAME_SUFFIX_PATTERN = Pattern.compile("[0-9]+");

    public static List<Job> parseJobInfos(List<String> jobInfos) {
        return jobInfos.stream().map(SlurmCmdExecuteResultParser::parseSQueueLine)
            .filter(Optional::isPresent)
            .map(Optional::get)
            .collect(Collectors.toList());
    }

    private static Optional<Job> parseSQueueLine(String line) {
        if (line.contains("JOBID")) {
            // 包含JOBID说明此行数据只是列名，不能解析
            return Optional.empty();
        }
        String[] columns = line.split("\\s+");
        if (columns.length < SQUEUE_MAX_COL) {
            throw new CmdExecuteResultParseException(
                String.format(Locale.ROOT, "squeue line is invalid, line is %s", line));
        }
        Job job = new Job();
        job.setId(columns[SQUEUE_ID_INDEX]);
        job.setName(columns[SQUEUE_NAME_INDEX]);
        job.setState(columns[SQUEUE_STATE_INDEX]);
        job.setUser(columns[SQUEUE_USER_INDEX]);
        job.setCpus(Integer.parseInt(columns[SQUEUE_CPUS_INDEX]));
        job.setMemory(Integer.parseInt(columns[SQUEUE_MEMORY_INDEX]));
        job.setTaskCount(0);
        return Optional.of(job);
    }

    public static List<Node> parseSInfoLine(String line) {
        String[] columns = line.split("\\s+");
        if (columns.length < SINFO_MAX_COL) {
            return Collections.emptyList();
        }
        NodeStatusEnum status = NodeStatusEnum.toNodeStatus(
            SlurmInstanceStatusEnum.value(columns[SINFO_STATE_INDEX]));
        String nodesInfo = columns[SINFO_NODELIST_INDEX];
        List<Node> nodes = new ArrayList<>();
        getNodeNameListFromNodesInfo(nodesInfo).forEach(name -> nodes.add(new Node(name, status)));
        return nodes;
    }

    static List<String> getNodeNameListFromNodesInfo(String nodesInfo) {
        List<String> names = new ArrayList<>();
        String[] nodeCells = nodesInfo.split(",");
        for (String cell : nodeCells) {
            if (cell.contains("-[")) {
                names.addAll(parseNodesRange(cell));
                continue;
            }
            names.add(cell);
        }
        return names;
    }

    static List<String> parseNodesRange(String nodes) {
        Matcher nodeNamePrefixMatcher = NODE_NAME_PREFIX_PATTERN.matcher(nodes);
        Matcher nodeNameSuffixSetMatcher = NODE_NAME_SUFFIX_SET_PATTERN.matcher(nodes);
        if (!nodeNamePrefixMatcher.find() || !nodeNameSuffixSetMatcher.find()) {
            return Collections.emptyList();
        }
        String nodeNamePrefix = nodeNamePrefixMatcher.group();
        String nodeNameSuffixSet = nodeNameSuffixSetMatcher.group();
        Matcher nodeNameSuffixMatcher = NODE_NAME_SUFFIX_PATTERN.matcher(nodeNameSuffixSet);

        nodeNameSuffixMatcher.find();
        int nodeSuffixHead = Integer.parseInt(nodeNameSuffixMatcher.group());
        int nodeSuffixLength = nodeNameSuffixMatcher.group().length();
        nodeNameSuffixMatcher.find();
        int nodeSuffixTail = Integer.parseInt(nodeNameSuffixMatcher.group());

        List<String> result = new ArrayList<>();
        for (int nodeSuffixNum = nodeSuffixHead; nodeSuffixNum <= nodeSuffixTail; nodeSuffixNum++) {
            String nodeSuffix = String.format("%7d", nodeSuffixNum).replace(" ", "0");
            String formattedNodeSuffix = nodeSuffix.substring(nodeSuffix.length() - nodeSuffixLength);
            result.add(nodeNamePrefix + "-" + formattedNodeSuffix);
        }
        return result;
    }
}